#!/usr/bin/python
# -*- coding: utf-8 -*-

__author__ = 'byrne_000'

import argparse
import boto
import boto.ec2
import boto.vpc
from config import config
import utils
from utils import N
from utils import tag


class NetDeployment:
    def __init__(self):
        self.route_table = None  #
        self.route_public = False
        self.route_table_subnets = []
        self.acl = None
        self.nets = [(None, False) for i in config.network.nets]


def query_net():
    net = NetDeployment()

    vpc_conn = boto.vpc.connect_to_region(config.REGION)
    rt_rst = vpc_conn.get_all_route_tables(filters=[('vpc_id', config.VPC_ID),
                                                    ('tag:Name', config.network.ROUTE_TABLE_NAME)])
    if len(rt_rst) == 1:
        net.route_table = rt_rst[0]
        for route in net.route_table.routes:
            if route.destination_cidr_block == config.network.CIDR_ANY and route.gateway_id == config.INTERNET_GATEWAY:
                net.route_public = True
                break
        for assoc in net.route_table.associations:
            net.route_table_subnets.append(assoc.id)

    nacl_rst = vpc_conn.get_all_network_acls(filters=[('vpc_id', config.VPC_ID), ('tag:Name', config.network.ACL)])
    if len(nacl_rst) == 1:
        net.acl = nacl_rst[0]

    for subnet in config.network.nets:
        _rst = vpc_conn.get_all_subnets(filters=[('vpcId', config.VPC_ID), ('cidrBlock', subnet['CIDR']),
                                                 ('availabilityZone', subnet['AZ'])])
        if len(_rst) == 1:
            net.nets[config.network.nets.index(subnet)] = (_rst[0], subnet['PUBLIC'])

    print('Subnets:')
    for subnet_tuple in net.nets:
        if subnet_tuple[0]:
            print('\t%s:\n\t\tAZ:%s\n\t\tCIDR:%s\n\t\tTags:%s\n\t\tIsPublic:%r' % (
                subnet_tuple[0].id, subnet_tuple[0].availability_zone,
                subnet_tuple[0].cidr_block, subnet_tuple[0].tags, subnet_tuple[1]))
    print('ACL:')
    if net.acl:
        print('\tid:%s\n\tTags:%s' % (net.acl.id, net.acl.tags))
        print('\tRules:')
        for entry in net.acl.network_acl_entries:
            print('\t\t#%s\n\t\t\taction:%s\n\t\t\tprotocol:%s\n\t\t\t%s\n\t\t\t%s\n\t\t\tegress:%s' % (
                entry.rule_number, entry.rule_action,
                entry.protocol, entry.port_range, entry.cidr_block, entry.egress))
    print('Route table:')
    if net.route_table:
        print('\tid:%s\n\tTags:%s' % (net.route_table.id, net.route_table.tags))
        print('\tRoutes:')
        for route in net.route_table.routes:
            print('\t\t(%s, %s)' % (route.gateway_id, route.destination_cidr_block))
    return net


def setup_net(net):
    vpc_conn = boto.vpc.connect_to_region(config.REGION)

    if not net.route_table:
        net.route_table = vpc_conn.create_route_table(config.VPC_ID)
        utils.tag(net.route_table, config.network.ROUTE_TABLE_NAME)

    if not net.route_public:
        if not vpc_conn.create_route(net.route_table.id, config.network.CIDR_ANY, config.INTERNET_GATEWAY):
            raise RuntimeError("Fail to setup internet gateway:" + config.INTERNET_GATEWAY)
        else:
            net.route_public = True

    # ACL
    if not net.acl:
        net.acl = vpc_conn.create_network_acl(config.VPC_ID)
        utils.tag(net.acl, config.network.ACL)

    # SUBNETS
    for subnet in net.nets:
        i = net.nets.index(subnet)
        netcfg = config.network.nets[i]
        if not subnet[0]:
            sub = vpc_conn.create_subnet(config.VPC_ID, netcfg['CIDR'], netcfg['AZ'])
            utils.tag(sub, netcfg['NAME'])
            assoc = vpc_conn.associate_route_table(net.route_table.id, sub.id)
            net.route_table_subnets.append(assoc)
            net.nets[i] = (sub, True if netcfg['PUBLIC'] else False)

            if not subnet[1]:  # if private subnet
                vpc_conn.associate_network_acl(network_acl_id=net.acl.id, subnet_id=net.nets[0][0].id)
                try:
                    vpc_conn.create_network_acl_entry(net.acl.id, rule_number=100,
                                                  protocol=6, rule_action='allow',
                                                  cidr_block=netcfg['CIDR'], egress=False,
                                                  port_range_from=0, port_range_to=65535)
                except boto.exception.EC2ResponseError as e:
                    if e.error_code == 'NetworkAclEntryAlreadyExists':
                        pass
                    else:
                        raise e


def remove_net(net):
    vpc_conn = boto.vpc.connect_to_region(config.REGION)
    if net.route_table:
        for assoc in net.route_table_subnets:
            vpc_conn.disassociate_route_table(assoc)
        if net.route_public:
            vpc_conn.delete_route(net.route_table.id, destination_cidr_block=config.network.CIDR_ANY)
        vpc_conn.delete_route_table(net.route_table.id)

    for subnet in net.nets:
        if subnet[1]:
            vpc_conn.disassociate_network_acl(subnet[0].id)

    if net.acl:
        vpc_conn.delete_network_acl(net.acl.id)

    for subnet in net.nets:
        if subnet[0]:
            try:
                vpc_conn.delete_subnet(subnet[0].id)
            except boto.exception.EC2ResponseError as e:
                print("Warning:" + e.message)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', metavar='command', choices=['setup', 'query', 'remove'],
                        default='setup', type=str, help='setup mode or remove mode')
    args = parser.parse_args()
    if args.mode == 'remove':
        net = query_net()
        remove_net(net)
        print('Updated information==============')
        query_net()

    elif args.mode == 'query':
        net = query_net()
    else:
        net = query_net()
        setup_net(net)
        print('Updated information==============')
        query_net()
